/*!
 * Copyright (c) 2012
 * Automation and Robotics Lab (LARA) at University of Brasilia (UnB)
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *     - Redistributions of source code must retain the above copyright
 *       notice, this list of conditions and the following disclaimer.
 *     - Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *     - Neither the name of the Automation and Robotics Lab (LARA) nor the
 *       names of its contributors may be used to endorse or promote products
 *       derived from this software without specific prior written permission.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License LGPL as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU Lesser General Public License LGPL for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License LGPL along with this program.
 * If not, see <http://www.gnu.org/licenses/>.
 *
 *****************************************************************
 *
 *   ROS stack name: schunk
 *   ROS package name: schunk_high
 *
 *   Authors: André A. Geraldes, email: andregeraldes@lara.unb.br
 *            Murilo M. Marinho, email: murilomarinho@lara.unb.br
 *
 *****************************************************************/

/******************************************************************************
                                   INCLUDES
 *****************************************************************************/

// General
#include <vector>
#include <string>

// Schunk_high
#include <schunk_high/cartesian_controller.h>

// Service Clients
#include <schunk_low/Init.h>
#include <schunk_low/Stop.h>
#include <schunk_low/Recover.h>
#include <schunk_low/Connect.h>
#include <schunk_low/Disconnect.h>
#include <schunk_low/ControlMode.h>
#include <schunk_low/GetPosVel.h>

/******************************************************************************
                                    CONSTRUCTORS AND DESTRUCTORS
 *****************************************************************************/

CartesianController::CartesianController(std::string controller_name, int controller_NID, int reference_size)
{
    // Allocate memory for the robot object
    robot_ = Robot_create();

    // Assignment of the defined callback queues to each node handler
    service_server_nodehandle_.setCallbackQueue(&service_server_queue_);
    service_client_nodehandle_.setCallbackQueue(&service_client_queue_);
    publisher_nodehandle_.setCallbackQueue(&publisher_queue_);
    references_nodehandle_.setCallbackQueue(&references_queue_);

    // Implementation of topics to publish
    publisher_velocity_ = publisher_nodehandle_.advertise<schunk_msgs::JointVelocities>("/schunk_low/control/velocities", 1);
    publisher_position_ = publisher_nodehandle_.advertise<schunk_msgs::JointPositions>("/schunk_low/control/positions", 1);

    // Implementation of service clients
    service_client_init_ = service_client_nodehandle_.serviceClient<schunk_low::Init>("/schunk_low/init");
    service_client_stop_ = service_client_nodehandle_.serviceClient<schunk_low::Stop>("/schunk_low/stop");
    service_client_recover_ = service_client_nodehandle_.serviceClient<schunk_low::Recover>("/schunk_low/recover");
    service_client_connect_ = service_client_nodehandle_.serviceClient<schunk_low::Connect>("/schunk_low/connect");
    service_client_disconnect_ = service_client_nodehandle_.serviceClient<schunk_low::Disconnect>("/schunk_low/disconnect");
    service_client_control_mode_ = service_client_nodehandle_.serviceClient<schunk_low::ControlMode>("/schunk_low/control_mode");
    service_client_get_pos_vel_ = service_client_nodehandle_.serviceClient<schunk_low::GetPosVel>("/schunk_low/get_pos_vel");

    // Setting status variables
    initialized_ = false;
    enabled_ = false;
    move_to_initial_position_ = false;

    robot_dh_loaded_ = false;
    control_parameters_loaded_ = false;
    initial_position_loaded_ = false;

    controller_name_ = controller_name;
    controller_NID_ = controller_NID;
    reference_size_ = reference_size;
}

CartesianController::~CartesianController()
{
}
/******************************************************************************
                                     PARAMETER LOADING FUNCTIONS
 *****************************************************************************/

bool CartesianController::GetRobotDH()
{
    bool error = false;

    // Read DH Matrix from the ROS Parameter Server
    XmlRpc::XmlRpcValue dh_xml_rpc;
    if (nodehandle_.hasParam("/schunk_high/schunk_dh"))
    {
        nodehandle_.getParam("/schunk_high/schunk_dh", dh_xml_rpc);
    }
    else
    {
        ROS_ERROR("Parameter schunk_dh not set, shutting down node...");
        nodehandle_.shutdown();
        error = true;
    }

    // Add links to the Robot object using the read matrix
    for (int i = 0; i < dh_xml_rpc.size(); i++)
    {
        if (dh_xml_rpc[i].size() == 4)
        {
            Robot_addLink(robot_, static_cast<double>(dh_xml_rpc[i][0]), static_cast<double>(dh_xml_rpc[i][1]), static_cast<double>(dh_xml_rpc[i][2]), static_cast<double>(dh_xml_rpc[i][3]));
        }
        else
        {
            ROS_ERROR("Incorrect size of DH matrix, shutting down node...");
            nodehandle_.shutdown();
            error = true;
        }
    }

    if (error) return false;

    publisher_position_message_.positions.resize(robot_->dofs);
    publisher_velocity_message_.velocities.resize(robot_->dofs);

    // Set the status variable
    robot_dh_loaded_ = true;
    return true;
}

bool CartesianController::GetRobotControlParameters()
{
    bool error = false;

    std::stringstream stringstream;
    std::string package_name = "/schunk_high/";

    // Sample Frequency
    stringstream << package_name << controller_name_ << "/sample_rate";
    if (nodehandle_.hasParam(stringstream.str()))
    {
        nodehandle_.getParam(stringstream.str(), sample_rate_);
    }
    else
    {
        ROS_WARN_STREAM("Parameter sample_rate not set, using standard frequency (25Hz)");
        sample_rate_ = 25;
    }
    stringstream.str("");

    // KP - Proportional Gains
    XmlRpc::XmlRpcValue kp_xml_rpc;
    stringstream << package_name << controller_name_ << "/kp";
    if (nodehandle_.hasParam(stringstream.str()))
    {
        nodehandle_.getParam(stringstream.str(), kp_xml_rpc);
    }
    else
    {
        ROS_ERROR("Parameter kp not set, shutting down node...");
        nodehandle_.shutdown();
        error = true;
    }
    stringstream.str("");

    // KI - Proportional Gains
    XmlRpc::XmlRpcValue ki_xml_rpc;
    stringstream << package_name << controller_name_ << "/ki";
    if (nodehandle_.hasParam(stringstream.str()))
    {
        nodehandle_.getParam(stringstream.str(), ki_xml_rpc);
    }
    else
    {
        ROS_ERROR("Parameter ki not set, shutting down node...");
        nodehandle_.shutdown();
        error = true;
    }
    stringstream.str("");

    // KD - Proportional Gains
    XmlRpc::XmlRpcValue kd_xml_rpc;
    stringstream << package_name << controller_name_ << "/kd";
    if (nodehandle_.hasParam(stringstream.str()))
    {
        nodehandle_.getParam(stringstream.str(), kd_xml_rpc);
    }
    else
    {
        ROS_ERROR("Parameter kd not set, shutting down node...");
        nodehandle_.shutdown();
        error = true;
    }
    stringstream.str("");

    // Saturation Velocity
    stringstream << package_name << controller_name_ << "/saturation_velocity";
    if (nodehandle_.hasParam(stringstream.str()))
    {
        nodehandle_.getParam(stringstream.str(), saturation_velocity_);
    }
    else
    {
        ROS_ERROR_STREAM("Parameter saturation_velocity not set, shutting down node...");
        nodehandle_.shutdown();
        error = true;
    }
    stringstream.str("");

    if ((kp_xml_rpc.size() != static_cast<int>(reference_size_)) || (ki_xml_rpc.size() != static_cast<int>(reference_size_)) || (kd_xml_rpc.size() != static_cast<int>(reference_size_)))
    {
        ROS_ERROR("Parameters size don't match the matrices order, shutting down node...");
        nodehandle_.shutdown();
        error = true;
    }

    if (error) return false;

    std::vector<double> kp;
    std::vector<double> ki;
    std::vector<double> kd;

    kp.resize(kp_xml_rpc.size());
    ki.resize(ki_xml_rpc.size());
    kd.resize(kd_xml_rpc.size());

    for (int i = 0; i < kp_xml_rpc.size(); i++) kp[i] = static_cast<double>(kp_xml_rpc[i]);
    for (int i = 0; i < ki_xml_rpc.size(); i++) ki[i] = static_cast<double>(ki_xml_rpc[i]);
    for (int i = 0; i < kd_xml_rpc.size(); i++) kd[i] = static_cast<double>(kd_xml_rpc[i]);

    // Init the Kinematic Variables using the read parameters
    if (!InitKinematicVariables(kp, ki, kd)) return false;

    control_parameters_loaded_ = true;

    return true;
}

bool CartesianController::GetRobotInitialPosition()
{
    bool error = false;

    std::stringstream stringstream;
    std::string package_name = "/schunk_high/";

    // Verify if an initial position is suposed to be used or not
    std::string move_to_initial_position_str;
    stringstream << package_name << controller_name_ << "/move_to_initial_position";
    if (nodehandle_.hasParam(stringstream.str()))
    {
        nodehandle_.getParam(stringstream.str(), move_to_initial_position_str);
    }
    else
    {
        ROS_ERROR("Parameter move_to_initial_position not set, shutting down node...");
        nodehandle_.shutdown();
        error = false;
    }
    if (!move_to_initial_position_str.compare("YES"))
    {
        move_to_initial_position_ = true;
    }
    else
    {
        move_to_initial_position_ = false;
    }
    stringstream.str("");

    // If there is an initial position, load it from the parameter server
    if (move_to_initial_position_)
    {
        XmlRpc::XmlRpcValue initial_position_xml_rpc;
        stringstream << package_name << controller_name_ << "/initial_position";
        if (nodehandle_.hasParam(stringstream.str()))
        {
            nodehandle_.getParam(stringstream.str(), initial_position_xml_rpc);
        }
        else
        {
            ROS_ERROR("Parameter /schunk_high/initial_position not set, shutting down node...");
            nodehandle_.shutdown();
            error = false;
        }
        stringstream.str("");

        if (!error)
        {
            initial_position_.resize(initial_position_xml_rpc.size());
            for (int i = 0; i < initial_position_xml_rpc.size(); i++)
            {
                initial_position_[i] = static_cast<double>(initial_position_xml_rpc[i]);
            }
        }
    }

    if (error) return false;

    initial_position_loaded_ = true;
    return true;
}

/******************************************************************************
                           SERVICE CLIENT WRAP FUNCTIONS
 *****************************************************************************/

bool CartesianController::InitSchunk()
{
    schunk_low::Init service_init_message;
    if (!service_client_init_.call(service_init_message))
    {
        ROS_ERROR_STREAM("Unable to init schunk_low control node");
        return false;
    }

    return true;
}

bool CartesianController::StopSchunk()
{
    schunk_low::Stop service_stop_message;
    if (!service_client_stop_.call(service_stop_message))
    {
        ROS_ERROR_STREAM("Unable to stop Schunk!");
        return false;
    }

    return true;
}

bool CartesianController::RecoverSchunk()
{
    schunk_low::Recover service_recover_message;
    if (!service_client_recover_.call(service_recover_message))
    {
        ROS_ERROR_STREAM("Unable to recover Schunk");
        return false;
    }

    return true;
}

bool CartesianController::ConnectToSchunk()
{
    schunk_low::Connect service_connect_message;
    service_connect_message.request.NID = controller_NID_;

    if (!service_client_connect_.call(service_connect_message))
    {
        ROS_ERROR_STREAM("Unable to connect to Schunk");
        return false;
    }

    return true;
}

bool CartesianController::DisconnectFromSchunk()
{
    schunk_low::Disconnect service_disconnect_message;
    if (!service_client_disconnect_.call(service_disconnect_message))
    {
        ROS_ERROR_STREAM("Unable to disconnect from Schunk");
        return false;
    }

    return true;
}

bool CartesianController::ChangeSchunkControlMode(int mode)
{
    schunk_low::ControlMode service_control_mode_message;
    service_control_mode_message.request.mode = mode;
    if (!service_client_control_mode_.call(service_control_mode_message))
    {
        ROS_ERROR_STREAM("Unable to change Shunk Control Mode");
        return false;
    }

    return true;
}

bool CartesianController::GetSchunkPositionsAndVelocitiesAndUpdateThetasAndDThetas()
{
    schunk_low::GetPosVel service_get_pos_vel_message;

    if (!service_client_get_pos_vel_.call(service_get_pos_vel_message))
    {
        ROS_ERROR_STREAM("Unable to get Schunk's positions and velocities, aborting");
        StopSchunk();
        return false;
    }

    else
    {
        for (int i = 0; i < robot_->dofs; i++){
            PGMATRIX_DATA(thetas_, i+1, 1) = service_get_pos_vel_message.response.positions[i];
            PGMATRIX_DATA(delta_thetas_, i+1, 1) = service_get_pos_vel_message.response.velocities[i];
        }
        return true;
    }
}

/******************************************************************************
                           PUBLISHER WRAP FUNCTIONS
 *****************************************************************************/

void CartesianController::SetSchunkVelocitiesAndSend()
{
    if (enabled_)
    {
        for (int i = 0; i < robot_->dofs; i++)
        {
            publisher_velocity_message_.velocities[i]=PGMATRIX_DATA(delta_thetas_, i+1, 1);

            if (publisher_velocity_message_.velocities[i] > saturation_velocity_)
            {
                publisher_velocity_message_.velocities[i] = saturation_velocity_;
                PGMATRIX_DATA(delta_thetas_, i+1, 1) = saturation_velocity_;
            }

            else if (publisher_velocity_message_.velocities[i] < -saturation_velocity_)
            {
                publisher_velocity_message_.velocities[i] = -saturation_velocity_;
                PGMATRIX_DATA(delta_thetas_, i+1, 1) = -saturation_velocity_;
            }
        }

        publisher_velocity_message_.header.stamp = ros::Time::now();
        publisher_velocity_message_.NID = controller_NID_;
        publisher_velocity_.publish(publisher_velocity_message_);
        publisher_queue_.callAvailable(ros::WallDuration(0));
    }
}

void CartesianController::SendJointPositionsToSchunk(std::vector<double> positions)
{
    if (enabled_)
    {
        for (int i = 0; i < robot_->dofs; i++)
        {
            publisher_position_message_.positions[i]=positions[i];
        }

        publisher_position_message_.header.stamp = ros::Time::now();
        publisher_position_message_.NID = controller_NID_;
        publisher_position_.publish(publisher_position_message_);
        publisher_queue_.callAvailable(ros::WallDuration(0));
    }
}

/******************************************************************************
                                      AUXILIAR FUNCTIONS
 *****************************************************************************/

bool CartesianController::InitKinematicVariables(std::vector<double> KP, std::vector<double> KI, std::vector<double> KD)
{
    if (KP.size() != reference_size_ || KI.size() != reference_size_ || KD.size() != reference_size_)
    {
        ROS_ERROR_STREAM("Cannot init kinematic variables - controller gains orders are wrong");
        return false;
    }

    // Initialize PID Errors
    kp_ = PGMATRIX_CREATE_DIAG_FROM_ARRAY(reference_size_, &KP[0]);
    ki_ = PGMATRIX_CREATE_DIAG_FROM_ARRAY(reference_size_, &KI[0]);
    kd_ = PGMATRIX_CREATE_DIAG_FROM_ARRAY(reference_size_, &KD[0]);

    // Initialize Thetas
    thetas_ = PGMATRIX_CREATE_ZEROES(robot_->dofs, 1);
    delta_thetas_ = PGMATRIX_CREATE_ZEROES(robot_->dofs, 1);

    // Reference Signal
    reference_state_variables_ = PGMATRIX_CREATE_ZEROES(reference_size_, 1);

    // Measured Values
    measured_state_variables_ = PGMATRIX_CREATE_ZEROES(reference_size_, 1);

    // Errors
    error_ = PGMATRIX_CREATE_ZEROES(reference_size_, 1);
    last_error_ = PGMATRIX_CREATE_ZEROES(reference_size_, 1);
    integral_error_ = PGMATRIX_CREATE_ZEROES(reference_size_, 1);

    // Task Jacobian's size depends on the task
    task_jacobian_ = NULL;

    // Auxiliar matrices
    aux_matrix_1_ = NULL;
    aux_matrix_2_ = NULL;
    aux_matrix_3_ = NULL;

    return true;
}

void CartesianController::ControlCleanUp()
{
    // Last Memory Free Procedure.
    // Jacobians
    PGMATRIX_FREE(task_jacobian_);
    // Errors
    PGMATRIX_FREE(error_);
    PGMATRIX_FREE(last_error_);
    PGMATRIX_FREE(integral_error_);
    // Var
    PGMATRIX_FREE(delta_thetas_);
    PGMATRIX_FREE(thetas_);
    PGMATRIX_FREE(reference_state_variables_);
    PGMATRIX_FREE(measured_state_variables_);
    // PID
    PGMATRIX_FREE(kp_);
    PGMATRIX_FREE(ki_);
    PGMATRIX_FREE(kd_);
    // Robot
    Robot_free(robot_);
}

bool CartesianController::MoveToInitialPosition(double tolerance)
{
    // Set the control mode to position
    if (!ChangeSchunkControlMode(CONTROL_MODE_POSITION))
    {
        DisconnectFromSchunk();
        return false;
    }

    // Wait for the schunk_low_control to subscribe to the positions topic
    ros::Rate poll_rate(1);
    while (publisher_position_.getNumSubscribers() == 0)
        poll_rate.sleep();

    // Move the robot to the initial position
    SendJointPositionsToSchunk(initial_position_);

    // Wait until the robot arrives at the initial position
    bool finished_moving = false;
    while (!finished_moving)
    {
        finished_moving = true;
        if (!GetSchunkPositionsAndVelocitiesAndUpdateThetasAndDThetas())
        {
            DisconnectFromSchunk();
            return false;
        }
        for (int i = 0; i < robot_->dofs; i++){
            if (PGMATRIX_DATA(thetas_, i+1, 1) > initial_position_[i] + tolerance){
                finished_moving = false;
            }
            if (PGMATRIX_DATA(thetas_, i+1, 1) < initial_position_[i] - tolerance){
                finished_moving = false;
            }
        }
        poll_rate.sleep();
    }

    return true;
}
